{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "e5b3c045-8b31-4481-af2c-f15b2004ce5f",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/ma-user/anaconda3/envs/MindSpore/lib/python3.9/site-packages/numpy/core/getlimits.py:499: UserWarning: The value of the smallest subnormal for <class 'numpy.float64'> type is zero.\n",
      "  setattr(self, word, getattr(machar, word).flat[0])\n",
      "/home/ma-user/anaconda3/envs/MindSpore/lib/python3.9/site-packages/numpy/core/getlimits.py:89: UserWarning: The value of the smallest subnormal for <class 'numpy.float64'> type is zero.\n",
      "  return self._float_to_str(self.smallest_subnormal)\n",
      "/home/ma-user/anaconda3/envs/MindSpore/lib/python3.9/site-packages/numpy/core/getlimits.py:499: UserWarning: The value of the smallest subnormal for <class 'numpy.float32'> type is zero.\n",
      "  setattr(self, word, getattr(machar, word).flat[0])\n",
      "/home/ma-user/anaconda3/envs/MindSpore/lib/python3.9/site-packages/numpy/core/getlimits.py:89: UserWarning: The value of the smallest subnormal for <class 'numpy.float32'> type is zero.\n",
      "  return self._float_to_str(self.smallest_subnormal)\n",
      "cannot found `mindformers.experimental`, please install dev version by\n",
      "`pip install git+https://gitee.com/mindspore/mindformers` \n",
      "or remove mindformers by \n",
      "`pip uninstall mindformers`\n",
      "Building prefix dict from the default dictionary ...\n",
      "Loading model from cache /tmp/jieba.cache\n",
      "Loading model cost 1.319 seconds.\n",
      "Prefix dict has been built successfully.\n"
     ]
    }
   ],
   "source": [
    "import argparse\n",
    "import os\n",
    "\n",
    "import numpy as np\n",
    "import mindspore as ms\n",
    "from mindspore import nn, ops\n",
    "from mindnlp.core import value_and_grad\n",
    "from mindspore.train import Model\n",
    "from mindspore import context, Tensor\n",
    "from mindnlp.transformers import BertGenerationConfig\n",
    "from mindnlp.transformers import BertGenerationDecoder\n",
    "from loaders.coco_full_loader import get_loader\n",
    "from mindnlp.core.optim import AdamW\n",
    "from tqdm.notebook import tqdm"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "bad79ead-724a-4181-b8b8-f6bd41c3115f",
   "metadata": {},
   "outputs": [],
   "source": [
    "def train_decoder(args):\n",
    "    # 设置运行环境\n",
    "    context.set_context(device_target=\"Ascend\")\n",
    "\n",
    "    # 初始化模型\n",
    "    if (not os.path.exists(f\"{args.saved_model_path}/decoder_model\")):\n",
    "        bert_config = BertGenerationConfig.from_pretrained(\"google/bert_for_seq_generation_L-24_bbc_encoder\")\n",
    "        bert_config.is_decoder = True\n",
    "        bert_config.add_cross_attention = True\n",
    "        bert_config.return_dict=True\n",
    "        bert_model = BertGenerationDecoder.from_pretrained(\"google/bert_for_seq_generation_L-24_bbc_encoder\",\n",
    "                                                           config=bert_config)\n",
    "    else:\n",
    "        bert_model = BertGenerationDecoder.from_pretrained(f\"{args.saved_model_path}/decoder_model\")\n",
    "\n",
    "    optimizer = AdamW(bert_model.trainable_params(),lr=args.lr,weight_decay=args.weight_decay)\n",
    "\n",
    "    # 定义前向网络\n",
    "    def forward_fn(input_ids, attention_mask, position_ids, clip_embeds, labels):\n",
    "        loss = bert_model(\n",
    "            input_ids=input_ids,\n",
    "            attention_mask=attention_mask,\n",
    "            position_ids=position_ids,\n",
    "            encoder_hidden_states=clip_embeds,\n",
    "            labels=labels\n",
    "        ).loss\n",
    "        return loss\n",
    "\n",
    "    # 定义梯度函数\n",
    "    grad_fn = value_and_grad(forward_fn,bert_model.trainable_params())\n",
    "\n",
    "    # 训练步骤\n",
    "    def train_step(input_ids, attention_mask, position_ids, clip_embeds, labels):\n",
    "        optimizer.zero_grad()\n",
    "        loss = grad_fn(input_ids, attention_mask, position_ids, clip_embeds, labels)\n",
    "        optimizer.step()\n",
    "        return loss\n",
    "\n",
    "    val_loss = evaluate(bert_model)\n",
    "    \n",
    "    # 训练循环\n",
    "    for epoch in range(args.num_epochs):\n",
    "        # 加载数据\n",
    "        train_dataset = get_loader(train=True, clip_backbone='ViT-B32')\n",
    "\n",
    "        bert_model.set_train()\n",
    "        total_loss = 0\n",
    "        steps=0\n",
    "\n",
    "        for step,batch in enumerate(tqdm(train_dataset.create_dict_iterator())):\n",
    "            input_ids = batch['input_ids']\n",
    "            attention_mask = batch['attention_mask']\n",
    "            labels = batch['label_ids']\n",
    "            clip_embeds = batch['clip_features']\n",
    "\n",
    "            # 生成position_ids\n",
    "            N, seq_length = input_ids.shape\n",
    "            # 先0~seq_len数据，然后使用python加轴，然后在维度0重复N次(这里会实际的创建数据，不是视图，基础深度学习框架也还是要去多学学)\n",
    "            position_ids = Tensor(np.arange(seq_length)[None].repeat(N, axis=0), ms.int32)\n",
    "\n",
    "            # 扩展clip_embeds\n",
    "            # ops.repeat_elements：这个操作会沿着指定的轴（axis参数）重复张量中的元素。具体来说：\n",
    "            clip_extended_embed = ops.repeat_elements(clip_embeds, rep=2, axis=1)\n",
    "            # ops.expand_dims：这个操作会在指定的位置插入一个新的维度。具体来说：\n",
    "            clip_extended_embed = ops.expand_dims(clip_extended_embed, 1)\n",
    "\n",
    "            loss = train_step(input_ids, attention_mask, position_ids,\n",
    "                              clip_extended_embed, labels)\n",
    "\n",
    "            total_loss += loss.asnumpy()\n",
    "            steps+=1\n",
    "\n",
    "        avg_loss = total_loss / steps\n",
    "        print(f'Epoch {epoch + 1}, Average Loss: {avg_loss}')\n",
    "\n",
    "        # 验证\n",
    "        val_loss = evaluate(bert_model)\n",
    "        print(f'Validation Loss: {val_loss}')\n",
    "\n",
    "        # 保存模型\n",
    "        bert_model.save_pretrained(f\"{args.saved_model_path}/decoder_model\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "248be719-e66c-4aeb-a58b-9434c7ddf101",
   "metadata": {},
   "outputs": [],
   "source": [
    "def evaluate(model):\n",
    "    model.set_train(False)\n",
    "    total_loss = 0\n",
    "    steps = 0\n",
    "    dataset = get_loader(train=False, clip_backbone='ViT-B32')\n",
    "\n",
    "    for step,batch in enumerate(tqdm(dataset.create_dict_iterator())):\n",
    "        input_ids = batch['input_ids']\n",
    "        attention_mask = batch['attention_mask']\n",
    "        labels = batch['label_ids']\n",
    "        clip_embeds = batch['clip_features']\n",
    "\n",
    "        N, seq_length = input_ids.shape\n",
    "        position_ids = Tensor(np.arange(seq_length)[None].repeat(N, axis=0), ms.int32)\n",
    "\n",
    "        clip_extended_embed = ops.repeat_elements(clip_embeds, rep=2, axis=1)\n",
    "        clip_extended_embed = ops.expand_dims(clip_extended_embed, 1)\n",
    "\n",
    "        loss = model(\n",
    "            input_ids=input_ids,\n",
    "            attention_mask=attention_mask,\n",
    "            position_ids=position_ids,\n",
    "            encoder_hidden_states=clip_extended_embed,\n",
    "            labels=labels\n",
    "        ).loss\n",
    "\n",
    "        total_loss += loss.asnumpy()\n",
    "        steps += 1\n",
    "\n",
    "    return total_loss / steps"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f4e34d0a-4240-4709-9f24-1b86dc3dd0de",
   "metadata": {
    "pycharm": {
     "is_executing": true
    }
   },
   "outputs": [],
   "source": [
    "import os\n",
    "import sys\n",
    "import argparse\n",
    "\n",
    "def get_args_in_notebook():\n",
    "    # 如果在 Jupyter Notebook 中运行，则直接定义参数\n",
    "    args = argparse.Namespace(\n",
    "        lr=1e-5,\n",
    "        weight_decay=1e-4,\n",
    "        num_epochs=24,\n",
    "        trained_path='./trained_models/COCO/'\n",
    "    )\n",
    "    return args\n",
    "\n",
    "if 'ipykernel' in sys.modules or 'IPython' in sys.modules:\n",
    "    # 检测是否在 Jupyter Notebook 或 IPython 环境中运行\n",
    "    args = get_args_in_notebook()\n",
    "else:\n",
    "    # 在命令行环境中正常解析参数\n",
    "    parser = argparse.ArgumentParser()\n",
    "    parser.add_argument('--lr', type=float, default=1e-5)\n",
    "    parser.add_argument('--weight_decay', type=float, default=1e-4)\n",
    "    parser.add_argument('--num_epochs', type=int, default=1)\n",
    "    parser.add_argument('--trained_path', type=str, default='./trained_models/COCO/')\n",
    "    args = parser.parse_args()\n",
    "\n",
    "# 动态设置保存模型路径\n",
    "args.saved_model_path = os.path.join(args.trained_path, 'ViT-B32/')\n",
    "\n",
    "# 创建目录（如果不存在）\n",
    "os.makedirs(args.saved_model_path, exist_ok=True)\n",
    "\n",
    "# 训练解码器\n",
    "train_decoder(args)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ca751c1c-f047-458e-814c-058f7dd5031f",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "916d2579-e609-42cc-9960-b0c1c8a6a745",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "MindSpore",
   "language": "python",
   "name": "mindspore"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
